package com.hr.huangzhk.handler;

import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * Created by huangzhk on 2016/10/20.
 */
@Component
public class MySocketHandler implements WebSocketHandler {

  public static final Map<Long, WebSocketSession> userSocketSessionMap;

  static {
    userSocketSessionMap = new HashMap<>();
  }

  @Override
  public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
    String userId = webSocketSession.getAttributes().get("userId").toString();
    if (StringUtils.isNotBlank(userId)) {
      // 将在线用户的webSocketSession添加到userSocketSessionMap中
      userSocketSessionMap.put(Long.parseLong(userId), webSocketSession);
      webSocketSession.sendMessage(new TextMessage("已成功监理websocket连接!"));
    }
  }

  @Override
  public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage)
      throws Exception {

  }

  @Override
  public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable)
      throws Exception {
    // 关闭session
    if (webSocketSession.isOpen()) {
      webSocketSession.close();
    }
    // 移除session
    removeUserSocketSession(webSocketSession);
  }

  @Override
  public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus)
      throws Exception {
    removeUserSocketSession(webSocketSession);
  }

  @Override
  public boolean supportsPartialMessages() {
    return false;
  }

  private void removeUserSocketSession(WebSocketSession webSocketSession) {
    Iterator<Map.Entry<Long, WebSocketSession>> it = userSocketSessionMap.entrySet().iterator();
    // 移除Socket会话
    while (it.hasNext()) {
      Map.Entry<Long, WebSocketSession> entry = it.next();
      if (entry.getValue().getId().equals(webSocketSession.getId())) {
        userSocketSessionMap.remove(entry.getKey());
        System.out.println("Socket会话已经移除:用户ID" + entry.getKey());
        break;
      }
    }
  }
}
